import torch
import torch.nn as nn
from efficientnet_pytorch import EfficientNet


def _make_encoder(backbone, features, use_pretrained, groups=1, expand=False, exportable=True):
    if backbone == "resnext101_wsl":
        pretrained = _make_pretrained_resnext101_wsl(use_pretrained)
        scratch = _make_scratch([256, 512, 1024, 2048], features, groups=groups, expand=expand)     # efficientnet_lite3  
    elif backbone == "efficientnet_lite3":
        pretrained = _make_pretrained_efficientnet_lite3(use_pretrained, exportable=exportable)
        scratch = _make_scratch([32, 48, 136, 384], features, groups=groups, expand=expand)  # efficientnet_lite3     
    else:
        print(f"Backbone '{backbone}' not implemented")
        assert False
        
    return pretrained, scratch


def _make_scratch(in_shape, out_shape, groups=1, expand=False):
    scratch = nn.Module()

    out_shape1 = out_shape
    out_shape2 = out_shape
    out_shape3 = out_shape
    out_shape4 = out_shape

    if expand==True:
        out_shape1 = out_shape
        out_shape2 = out_shape*2
        out_shape3 = out_shape*4
        out_shape4 = out_shape*8

    scratch.layer1_rn = nn.Sequential(
        nn.ReflectionPad2d(1),
        nn.Conv2d(
        in_shape[0], out_shape1, kernel_size=3, stride=1, padding=0, bias=False, groups=groups
    ))
    scratch.layer2_rn = nn.Sequential(
        nn.ReflectionPad2d(1),
        nn.Conv2d(
        in_shape[1], out_shape2, kernel_size=3, stride=1, padding=0, bias=False, groups=groups
    ))
    scratch.layer3_rn = nn.Sequential(
        nn.ReflectionPad2d(1),
        nn.Conv2d(
        in_shape[2], out_shape3, kernel_size=3, stride=1, padding=0, bias=False, groups=groups
    ))
    scratch.layer4_rn = nn.Sequential(
        nn.ReflectionPad2d(1),
        nn.Conv2d(
        in_shape[3], out_shape4, kernel_size=3, stride=1, padding=0, bias=False, groups=groups
    ))

    return scratch


def _make_pretrained_efficientnet_lite3(use_pretrained, exportable=False):
    efficientnet = torch.hub.load(
        "/home/ubuntu/DCM-dehaze-main/rwightman/gen-efficientnet-pytorch-master",
        "tf_efficientnet_lite3",
        # path='/home/ubuntu/D4-main_test/rwightman/gen-efficientnet-pytorch/tf_efficientnet_lite3.pth',
        source="local",
        pretrained=use_pretrained,
        exportable=exportable
    )
    
    return _make_efficientnet_backbone(efficientnet)

#efficientnet="/home/ubuntu/UR2P-Dehaze/model.pth"
# def _make_pretrained_efficientnet_lite3(use_pretrained, exportable=False):

#     # efficientnet = timm.create_model(
#     #     "tf_efficientnet_lite3", pretrained=use_pretrained, exportable=exportable
#     # )
    
#     # weight_path = "/home/ubuntu/D4-main_test/tf_efficientnet_lite3.pth"  # 权重文件的路径
#     # state_dict = torch.load(weight_path)
#     # efficientnet.load_state_dict(state_dict)
    
#     efficientnet = EfficientNet.from_name('efficientnet-b3')
    
#     weight_path = "/home/ubuntu/D4-main_test/tf_efficientnet_lite3.pth"  # 权重文件的路径
#     state_dict = torch.load(weight_path)
#     efficientnet.load_state_dict(state_dict)

#     return _make_efficientnet_backbone(efficientnet)


def _make_efficientnet_backbone(effnet):
    pretrained = nn.Module()

    pretrained.layer1 = nn.Sequential(
        effnet.conv_stem, effnet.bn1, effnet.act1, *effnet.blocks[0:2]
    )
    pretrained.layer2 = nn.Sequential(*effnet.blocks[2:3])
    pretrained.layer3 = nn.Sequential(*effnet.blocks[3:5])
    pretrained.layer4 = nn.Sequential(*effnet.blocks[5:9])

    return pretrained
    

def _make_resnet_backbone(resnet):
    pretrained = nn.Module()
    pretrained.layer1 = nn.Sequential(
        resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1
    )

    pretrained.layer2 = resnet.layer2
    pretrained.layer3 = resnet.layer3
    pretrained.layer4 = resnet.layer4

    return pretrained


def _make_pretrained_resnext101_wsl(use_pretrained):
    resnet = torch.hub.load("facebookresearch/WSL-Images", "resnext101_32x8d_wsl")
    return _make_resnet_backbone(resnet)
# def _make_pretrained_resnext101_wsl(use_pretrained):
#     if use_pretrained:
#         weight_path = "/home/ubuntu/D4-main_test/resnext101_32x8d_wsl.pth"  # 权重文件的本地路径
#         state_dict = torch.load(weight_path)
#         resnet = torch.hub.load(
#             "facebookresearch/WSL-Images",
#             "resnext101_32x8d_wsl",
#             state_dict=state_dict,
#             pretrained=False  # 设置为False，不从 Torch Hub 下载权重
#         )
#     else:
#         resnet = torch.hub.load("facebookresearch/WSL-Images", "resnext101_32x8d_wsl")
    
#     return _make_resnet_backbone(resnet)



class Interpolate(nn.Module):
    """Interpolation module.
    """

    def __init__(self, scale_factor, mode):
        """Init.

        Args:
            scale_factor (float): scaling
            mode (str): interpolation mode
        """
        super(Interpolate, self).__init__()

        self.interp = nn.functional.interpolate
        self.scale_factor = scale_factor
        self.mode = mode

    def forward(self, x):
        """Forward pass.

        Args:
            x (tensor): input

        Returns:
            tensor: interpolated data
        """

        x = self.interp(
            x, scale_factor=self.scale_factor, mode=self.mode, align_corners=False
        )

        return x


class ResidualConvUnit(nn.Module):
    """Residual convolution module.
    """

    def __init__(self, features):
        """Init.

        Args:
            features (int): number of features
        """
        super().__init__()

        self.conv1 = nn.Conv2d(
            features, features, kernel_size=3, stride=1, padding=1, bias=True
        )

        self.conv2 = nn.Conv2d(
            features, features, kernel_size=3, stride=1, padding=1, bias=True
        )

        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        """Forward pass.

        Args:
            x (tensor): input

        Returns:
            tensor: output
        """
        out = self.relu(x)
        out = self.conv1(out)
        out = self.relu(out)
        out = self.conv2(out)

        return out + x


class FeatureFusionBlock(nn.Module):
    """Feature fusion block.
    """

    def __init__(self, features):
        """Init.

        Args:
            features (int): number of features
        """
        super(FeatureFusionBlock, self).__init__()

        self.resConfUnit1 = ResidualConvUnit(features)
        self.resConfUnit2 = ResidualConvUnit(features)

    def forward(self, *xs):
        """Forward pass.

        Returns:
            tensor: output
        """
        output = xs[0]

        if len(xs) == 2:
            output += self.resConfUnit1(xs[1])

        output = self.resConfUnit2(output)

        output = nn.functional.interpolate(
            output, scale_factor=2, mode="bilinear", align_corners=True
        )

        return output




class ResidualConvUnit_custom(nn.Module):
    """Residual convolution module.
    """

    def __init__(self, features, activation, bn):
        """Init.

        Args:
            features (int): number of features
        """
        super().__init__()

        self.bn = bn

        self.groups=1

        self.conv1 = nn.Conv2d(
            features, features, kernel_size=3, stride=1, padding=0, bias=True, groups=self.groups
        )
        
        self.conv2 = nn.Conv2d(
            features, features, kernel_size=3, stride=1, padding=0, bias=True, groups=self.groups
        )

        self.pad1 = nn.ReflectionPad2d(1)

        if self.bn==True:
            self.bn1 = nn.BatchNorm2d(features)
            self.bn2 = nn.BatchNorm2d(features)

        self.activation = activation

        self.skip_add = nn.quantized.FloatFunctional()

    def forward(self, x):
        """Forward pass.

        Args:
            x (tensor): input

        Returns:
            tensor: output
        """
        
        out = self.activation(x)
        out = self.pad1(out)
        out = self.conv1(out)
        if self.bn==True:
            out = self.bn1(out)
       
        out = self.activation(out)
        out = self.pad1(out)
        out = self.conv2(out)
        if self.bn==True:
            out = self.bn2(out)

        if self.groups > 1:
            out = self.conv_merge(out)

        return self.skip_add.add(out, x)



class FeatureFusionBlock_custom(nn.Module):
    """Feature fusion block.
    """

    def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True):
        """Init.

        Args:
            features (int): number of features
        """
        super(FeatureFusionBlock_custom, self).__init__()

        self.deconv = deconv
        self.align_corners = align_corners


                # 添加注意力模块
        self.attention_conv1 = nn.Conv2d(features, features // 16, kernel_size=1, stride=1, padding=0)
        self.attention_activation = nn.ReLU(inplace=True)
        self.attention_conv2 = nn.Conv2d(features // 16, 1, kernel_size=1, stride=1, padding=0)
        self.attention_sigmoid = nn.Sigmoid()

        self.expand = expand
        out_features = features
        if self.expand:
            out_features = features // 2
        
        self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)

        self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn)
        self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn)
        
        self.skip_add = nn.quantized.FloatFunctional()

    def forward(self, *xs):
        """Forward pass.

        Returns:
            tensor: output
        """
        output = xs[0]

        if len(xs) == 2:
            res = self.resConfUnit1(xs[1])
            output = self.skip_add.add(output, res)

        # # 注意力模块
            
        attention = self.attention_conv1(output)
        attention = self.attention_activation(attention)
        attention = self.attention_conv2(attention)
        attention = self.attention_sigmoid(attention)

        output = self.resConfUnit2(output)
        
        output = output * attention


        

        output = nn.functional.interpolate(
            output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
        )

        output = self.out_conv(output)

        return output

